# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import ujson as json

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.bithumb.kr_rest.product import BithumbProduct
from coin.exchange.bithumb.book_builder import BookBuilder


def parse_trade_side(side_str):
  if side_str == 'ask':
    side = TradeSide.TRADE_SELL_SIDE
  elif side_str == 'bid':
    side = TradeSide.TRADE_BUY_SIDE
  elif side_str == 'dn':
    side = TradeSide.TRADE_SELL_SIDE
  elif side_str == 'up':
    side = TradeSide.TRADE_BUY_SIDE
  else:
    side = TradeSide.TRADE_NO_SIDE
  return side


def parse_trade_side_2(side_str):
  if side_str == '1':
    side = TradeSide.TRADE_SELL_SIDE
  elif side_str == '2':
    side = TradeSide.TRADE_BUY_SIDE
  else:
    side = TradeSide.TRADE_NO_SIDE
  return side


class BithumbHandler(WebsocketHandler):
  def __init__(self, *, feed_checker_config=None):
    super().__init__(feed_checker_config=feed_checker_config)
    self._last_trades = {}

  def on_book(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, (list, dict)):
      data = json.loads(data.tobytes())

    if data['status'] != '0000':
      return

    header = data['header']
    assert header['service'] == 'orderbook', data
    product = BithumbProduct.FromStrNativeProduct(header['currency'])

    if product.native_symbol not in self._books:
      self._books[product.native_symbol] = BookBuilder(
          product.native_symbol, book_builder_observer=self._book_builder_observer)
    book = self._books[product.native_symbol]

    book.update(data['data']['bids'], data['data']['asks'])
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', product.native_symbol, book)
    return book

  def on_trade(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict):
      message = json.loads(message.tobytes())

    assert message['header']['service'] == 'transaction', message
    symbol = message['header']['currency']
    product = BithumbProduct.FromStrNativeProduct(symbol)
    if not self.is_subscribed('trade', symbol):
      return
    # Bithumb trade messages are given in reversed order.
    for trade_msg in reversed(message['data']):
      price = float(trade_msg['price'])
      qty = float(trade_msg['units_traded'])
      side = parse_trade_side(trade_msg['type'])
      cont_no = int(trade_msg['cont_no'])
      if cont_no <= self._last_trades.get(symbol, -1):
        continue
      trade_event = create_market_trade_event(record.timestamp,
                                              symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)
      self.publish_event('trade', symbol, trade_event)

    if len(message['data']) > 0:
      last_trade = message['data'][0]
      self._last_trades[symbol] = int(last_trade['cont_no'])
